def depth_to_space(tensor, scale_factor):  # scale_factor = 2
    num, ch, height, width = tensor.shape  # n 4 128 128
    if ch % (scale_factor * scale_factor) != 0:
        raise ValueError('channel of tensor must be divisible by '
                         '(scale_factor * scale_factor).必须是整数倍.')

    new_ch = ch // (scale_factor * scale_factor)  # 1
    new_height = height * scale_factor  # 256
    new_width = width * scale_factor  # 256

    tensor = tensor.reshape(
        [num, new_ch, scale_factor, scale_factor, height, width])  # n 1 2 2 128 128
    tensor = tensor.permute([0, 1, 4, 2, 5, 3])  # n 1 128 2 128 2
    tensor = tensor.reshape([num, new_ch, new_height, new_width])  # n 1 256 256
    return tensor


if __name__ == '__main__':
    from s2d import space_to_depth
    import torch

    data = torch.range(0, 71).resize(1, 4, 6, 3).permute(0, 3, 1, 2)
    print(data)
    print(depth_to_space(space_to_depth(data, 2).reshape(1, 12, 2, 3), 2))